#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import sys
import rospy
import numpy as np
import tf
from sensor_msgs.msg import JointState
from geometry_msgs.msg import Pose, Point
from std_msgs.msg import Float64
import moveit_msgs.msg
import geometry_msgs.msg
import tf.transformations
import moveit_commander
from moveit_commander import MoveGroupCommander
from copy import deepcopy

# 全局变量存储接收到的数据
joint_state_data = None
result2_data = None

# 创建全局发布者
tf_viewer = None

def joint_states_callback(data):
    global joint_state_data
    joint_state_data = data

def result2_callback(data):
    global result2_data
    result2_data = data

def process_data(event):
    global joint_state_data, result2_data, tf_viewer
    if joint_state_data is not None and result2_data is not None:
            
            # 获取前五个关节的位置并调整theta4
            joint_positions = [joint_state_data.position[i] if i < len(joint_state_data.position) else 0 for i in range(5)]
            # 打印前五个关节的位置（弧度）
            # rospy.loginfo('joint_states: %s', joint_positions)
            # 打印前五个关节的位置(角度)
            joint_positions_degrees = [np.degrees(theta) for theta in joint_positions]
            rospy.loginfo('joint_states (degrees): %s', joint_positions_degrees)
            # 将所有角度值改为负的
            joint_positions = [-theta for theta in joint_positions]
            joint_positions[1] -= np.pi/2  # 调整theta2
            joint_positions[3] += np.pi/2  # 调整theta4

            # # 打印前五个关节的位置（弧度）
            # rospy.loginfo('joint_states: %s', joint_positions)
            # 打印/result2_topic接收到的xyz
            rospy.loginfo('result2_topic xyz: %s', (result2_data.x, result2_data.y, result2_data.z))

            # 定义theta变量
            theta1, theta2, theta3, theta4, theta5 = joint_positions
            theta6 = 0  # 定义theta6为0

            # 打印调整后的五个关节位置（角度）
            joint_positions_degrees = [np.degrees(theta) for theta in joint_positions]
            # rospy.loginfo('Adjusted joint_states (degrees): %s', joint_positions_degrees)

            cos = np.cos
            sin = np.sin

            TS = np.array([
                [cos(theta6)*(cos(theta1)*cos(theta5) - cos(theta2 + theta3 + theta4)*sin(theta1)*sin(theta5)) - sin(theta6)*(cos(theta1)*sin(theta5) + cos(theta2 + theta3 + theta4)*cos(theta5)*sin(theta1)), sin(theta2 + theta3 + theta4)*sin(theta1), -sin(theta6)*(cos(theta1)*cos(theta5) - cos(theta2 + theta3 + theta4)*sin(theta1)*sin(theta5)) - cos(theta6)*(cos(theta1)*sin(theta5) + cos(theta2 + theta3 + theta4)*cos(theta5)*sin(theta1)), 120*sin(theta1)*sin(theta2)*sin(theta3) - 94*cos(theta1)*sin(theta5) - 120*cos(theta2)*sin(theta1) - 94*cos(theta2 + theta3 + theta4)*cos(theta5)*sin(theta1) - (117*cos(theta2 + theta3)*cos(theta4)*sin(theta1))/4 - (177*cos(theta2 + theta3)*sin(theta1)*sin(theta4))/5 - (177*sin(theta2 + theta3)*cos(theta4)*sin(theta1))/5 + (117*sin(theta2 + theta3)*sin(theta1)*sin(theta4))/4 - 120*cos(theta2)*cos(theta3)*sin(theta1)],
                [cos(theta6)*(cos(theta5)*sin(theta1) + cos(theta2 + theta3 + theta4)*cos(theta1)*sin(theta5)) - sin(theta6)*(sin(theta1)*sin(theta5) - cos(theta2 + theta3 + theta4)*cos(theta1)*cos(theta5)), -sin(theta2 + theta3 + theta4)*cos(theta1), -sin(theta6)*(cos(theta5)*sin(theta1) + cos(theta2 + theta3 + theta4)*cos(theta1)*sin(theta5)) - cos(theta6)*(sin(theta1)*sin(theta5) - cos(theta2 + theta3 + theta4)*cos(theta1)*cos(theta5)), 120*cos(theta1)*cos(theta2) - 94*sin(theta1)*sin(theta5) - 120*cos(theta1)*sin(theta2)*sin(theta3) + 94*cos(theta2 + theta3 + theta4)*cos(theta1)*cos(theta5) + (117*cos(theta2 + theta3)*cos(theta1)*cos(theta4))/4 + (177*cos(theta2 + theta3)*cos(theta1)*sin(theta4))/5 + (177*sin(theta2 + theta3)*cos(theta1)*cos(theta4))/5 - (117*sin(theta2 + theta3)*cos(theta1)*sin(theta4))/4 + 120*cos(theta1)*cos(theta2)*cos(theta3) - 211/2],
                [-sin(theta2 + theta3 + theta4)*sin(theta5 + theta6), -cos(theta2 + theta3 + theta4), -sin(theta2 + theta3 + theta4)*cos(theta5 + theta6), (177*cos(theta2 + theta3)*cos(theta4))/5 - 120*sin(theta2) - (177*sin(theta2 + theta3)*sin(theta4))/5 - cos(theta5)*(94*cos(theta2 + theta3)*sin(theta4) + 94*sin(theta2 + theta3)*cos(theta4)) - 120*sin(theta2 + theta3) - (117*cos(theta2 + theta3)*sin(theta4))/4 - (117*sin(theta2 + theta3)*cos(theta4))/4 + 2301/20],
                [0, 0, 0, 1]
            ])

            # rospy.loginfo("TS:\n%s", TS)

            # 定义M矩阵
            M = np.eye(4)
            M[0, 3] = result2_data.x
            M[1, 3] = result2_data.y
            M[2, 3] = result2_data.z

            # 其他矩阵转换
            T67_1 = np.array([[0, -1, 0, 0],
                              [0, 0, -1, 0],
                              [1, 0, 0, 0],
                              [0, 0, 0, 1]])
            Tcs1 = np.array([[1, 0, 0, 51],
                             [0, 1, 0, 0],
                             [0, 0, 1, -50.82],
                             [0, 0, 0, 1]])
            Tcs2 = np.array([[0, 1, 0, 0],
                             [-1, 0, 0, 0],
                             [0, 0, 1, 0],
                             [0, 0, 0, 1]])
            T12 = np.array([[1, 0, 0, 0],
                             [0, 0, -1, 0],
                             [0, 1, 0, 0],
                             [0, 0, 0, 1]])

            # 结果矩阵计算 camera乘节点数据
            result_matrix = TS @ M @ T12
            result_matrix[:3, 3] /= 1000  # 调整尺度

            # 打印和发布结果
            # rospy.loginfo('TS:\n %s', TS)
            # rospy.loginfo('Result Matrix:\n %s', result_matrix)

            # 提取位置和旋转矩阵
            position = result_matrix[:3, 3]
            rotation_matrix = result_matrix[:3, :3]
            
            # 将旋转矩阵转换为四元数
            quaternion = tf.transformations.quaternion_from_matrix(result_matrix)
            
            # # 打印位置和四元数
            # print("Extracted Position:")
            # print(f"x: {position[0]}, y: {position[1]}, z: {position[2]}")
            # print("Extracted Quaternion:")
            # print(f"x: {quaternion[0]}, y: {quaternion[1]}, z: {quaternion[2]}, w: {quaternion[3]}")
            
            # 设置数据处理标志
            data_processed = True

            # 通过MoveIt控制机械臂
            move_arm_to_position(position, quaternion)

def move_arm_to_position(position, quaternion):
    global arm
    # 设置目标位姿
    target_pose = Pose()
    target_pose.position.x = position[0]
    target_pose.position.y = position[1]
    target_pose.position.z = position[2]
    target_pose.orientation.x = quaternion[0]
    target_pose.orientation.y = quaternion[1]
    target_pose.orientation.z = quaternion[2]
    target_pose.orientation.w = quaternion[3]

    # 控制机械臂运动到目标位姿
    arm.set_pose_target(target_pose)
    success = arm.go()
    rospy.sleep(1)

    if success:
        print("规划成功")
    else:
        print("未规划")
def listener():
    global arm
    rospy.init_node('multi_topic_listener', anonymous=True)  # 保留这个初始化节点的调用

    # 初始化move_group的API
    moveit_commander.roscpp_initialize(sys.argv)
    
    # 初始化需要使用move group控制的机械臂中的arm group
    arm = MoveGroupCommander('manipulator')

    # 当运动规划失败后，允许重新规划
    arm.allow_replanning(True)

    # 设置目标位置所使用的参考坐标系
    arm.set_pose_reference_frame('base_link')

    # 设置位置(单位：米)和姿态（单位：弧度）的允许误差
    arm.set_goal_position_tolerance(0.0001)
    arm.set_goal_orientation_tolerance(0.0001)

    # 设置允许的最大速度和加速度
    arm.set_max_acceleration_scaling_factor(0.1)
    arm.set_max_velocity_scaling_factor(0.1)

    # 控制机械臂先回到初始化位置
    arm.set_named_target('home')
    arm.go()
    rospy.sleep(1)

    global tf_viewer
    tf_viewer = rospy.Publisher('result_data', Float64, queue_size=10)
    rospy.Subscriber("/joint_states", JointState, joint_states_callback)
    # rospy.Subscriber("/result2_topic", Point, result2_callback)

    ## 模拟接收到 result2_topic 消息
    global result2_data
    result2_data = Point(0, 40.82, 101)

    rospy.Timer(rospy.Duration(3), process_data)
    rospy.spin()


if __name__ == '__main__':
    try:
        moveit_commander.roscpp_initialize(sys.argv)
        listener()
    except rospy.ROSInterruptException:
        pass
